import math
import random

from torch import Tensor
from torchvision.transforms import functional as F
from torchvision.transforms.functional import F_t


# 随机 裁剪图片为正方形， 最端边做为 正方形长
def to_square(src: Tensor) -> Tensor:
    src_h, src_w = src.shape[-2:]

    if src_h != src_w:
        if src_h > src_w:
            top = random.randint(0, src_h - src_w + 1)
            left = 0
            crop_t = src_w
        else:
            top = 0
            left = random.randint(0, src_w - src_h + 1)
            crop_t = src_h
        dst = F_t.crop(src, top, left, crop_t, crop_t)
    else:
        dst = src

    return dst


def rotate_by_factor(src: Tensor, angle_deg: float) -> Tensor:
    if angle_deg != 0:
        dst = F.rotate(src, angle_deg, F.InterpolationMode.BILINEAR)
    else:
        dst = src

    return dst


def rotate_square(src: Tensor, angle_deg: float) -> Tensor:
    src_size, src_w = src.shape[-2:]
    assert src_size == src_w

    dst = rotate_by_factor(src, angle_deg)

    angle_deg = angle_deg % 90
    angle_rad = math.radians(angle_deg)

    div_factor = math.sin(angle_rad) + math.cos(angle_rad)
    if div_factor != 1.0:
        crop_t = src_size / div_factor
        dst = F.center_crop(dst, crop_t)

    return dst


# img = Image.open(test_image_path).convert("RGB")
def img_rotate(img, angle_deg):
    # 图片转 张量 并 归一化处理 [0, 1]
    img_tensor = F.to_tensor(img)

    pil_tensor = to_square(img_tensor)

    # 旋转后的图片
    dst = rotate_square(pil_tensor, angle_deg)

    return dst
